from __future__ import annotations

import json
from typing import Any

from pydantic import BaseModel, ValidationError

from aws_lambda_powertools import Logger, Metrics, Tracer
from aws_lambda_powertools.utilities.batch import BatchProcessor, EventType, process_partial_response
from aws_lambda_powertools.utilities.typing import LambdaContext

logger = Logger()
tracer = Tracer()
metrics = Metrics()

# Initialize batch processor for SQS
processor = BatchProcessor(event_type=EventType.SQS)


class WorkerMessage(BaseModel):
    task_id: str
    task_type: str
    payload: dict


@tracer.capture_method
def record_handler(record):
    """Process individual SQS record"""
    try:
        # Parse and validate message
        message_data = json.loads(record.body)
        worker_message = WorkerMessage(**message_data)

        logger.info("Processing task", extra={"task_id": worker_message.task_id, "task_type": worker_message.task_type})

        # Simulate work based on task type
        if worker_message.task_type == "email":
            # Process email task
            logger.info("Sending email", extra={"task_id": worker_message.task_id})
        elif worker_message.task_type == "report":
            # Process report task
            logger.info("Generating report", extra={"task_id": worker_message.task_id})
        else:
            logger.warning("Unknown task type", extra={"task_type": worker_message.task_type})

        metrics.add_metric(name="TaskProcessed", unit="Count", value=1)
        metrics.add_metadata(key="task_type", value=worker_message.task_type)

        return {"status": "success", "task_id": worker_message.task_id}

    except ValidationError as e:
        logger.error("Invalid message format", extra={"error": str(e)})
        metrics.add_metric(name="TaskFailed", unit="Count", value=1)
        raise
    except Exception as e:
        logger.error("Task processing failed", extra={"error": str(e)})
        metrics.add_metric(name="TaskFailed", unit="Count", value=1)
        raise


@logger.inject_lambda_context
@tracer.capture_lambda_handler
@metrics.log_metrics
def lambda_handler(event: dict[str, Any], context: LambdaContext):
    """Process SQS messages using BatchProcessor"""

    return process_partial_response(
        event=event,
        record_handler=record_handler,
        processor=processor,
        context=context,
    )
